#python -m lilab.OpenLabCluster_train.a6b_clippredpkl_2_cluster400p clipredpkl dir_400p
#%%

import pickle
import os.path as osp
import numpy as np
import os
import tqdm
from lilab.comm_signal.balanceSample import balanceSample
import ffmpegcv
import mmcv
from lilab.openlabcluster_postprocess.s1_merge_3_file import get_assert_1_file
import argparse


def create_400p_clusters(clipredpkl, dir_400p):
    clippreddata = pickle.load(open(clipredpkl, "rb"))
    df_clipNames = clippreddata["df_clipNames"]
    cluster_labels = clippreddata["cluster_labels"]
    ntwin = clippreddata["ntwin"]
    video = df_clipNames[['vnake', 'isBlack', 'startFrame']].values
    isBlackFirst = video[:,1].astype(bool)
    videoB = video[isBlackFirst]
    cluster_labelsB = cluster_labels[isBlackFirst]
    _, indchoose, _ = balanceSample(cluster_labelsB, by='number', by_value=60)
    videoB_c = videoB[indchoose]
    cluster_labelsB_c = cluster_labelsB[indchoose]

    cluster_u = np.unique(cluster_labelsB_c)
    outvid_dir = osp.join(osp.dirname(clipredpkl), "400p_clusters")
    os.makedirs(outvid_dir, exist_ok=True)

    vin_file_dict = {v:osp.join(dir_400p,  v+'_1_400p.mp4') for v in set(videoB_c[:,0])}
    for f in vin_file_dict.values():
        if not osp.exists(f):
            print('NOT FOUND: ', osp.basename(f))
    assert all([osp.exists(f) for f in vin_file_dict.values()])
    vin_dict = {v:mmcv.VideoReader(f) for v,f in vin_file_dict.items()}
    vidout_l = [ffmpegcv.VideoWriter(osp.join(outvid_dir, f'clu_{cluster_u_i}.mp4'), fps=30)
                for cluster_u_i in cluster_u]
    for cluster_i, vidout_i in zip(tqdm.tqdm(cluster_u), vidout_l):
        ind_i = np.where(cluster_labelsB_c == cluster_i)[0]
        videoB_c_clusteri = videoB_c[ind_i][:,[0,2]]
        for v, start_frame in videoB_c_clusteri:
            start_frame = int(start_frame)
            for frame in range(start_frame, start_frame+ntwin):
                if  frame>=len(vin_dict[v]): continue
                vidout_i.write(vin_dict[v][frame])
    
    for vidout_i in vidout_l:
        vidout_i.release()

    return outvid_dir    


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('clipredpkl', type=str, help='clippredpkl')
    parser.add_argument('dir_400p', type=str, help='dir_400p')
    args = parser.parse_args()
    outvid_dir = create_400p_clusters(args.clipredpkl, args.dir_400p)
    print(f'Save video to {outvid_dir}')
